{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pseudopotential Test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_xyz\n",
    "from ase.calculators.vasp import Vasp\n",
    "import shutil\n",
    "import os\n",
    "\n",
    "frames = read_xyz('test.xyz')\n",
    "for i in range(len(frames)):\n",
    "    atoms = frames[i]\n",
    "    dir = 'Your_Path'\n",
    "    if os.path.exists(dir):\n",
    "        shutil.rmtree(dir)\n",
    "    calc = Vasp(directory=dir,\n",
    "                encut=500,                  # ENCUT\n",
    "                ediff=1e-6,                 # EDIFF\n",
    "                kspacing=0.15,              # KSPACING\n",
    "                sigma=0.1,                  # SIGMA\n",
    "                ismear=1,                   # ISMEAR\n",
    "                gga='PE',                   # GGA\n",
    "                prec='Accurate',            # PREC\n",
    "                lasph=True,                 # LASPH\n",
    "                nelmin=4,                   # NELMIN\n",
    "                nelm=200,                   # NELM\n",
    "                kpar=4,                     # KPAR\n",
    "                ncore=20,                   # NCORE\n",
    "                lwave=False)                # LWAVE      \n",
    "    calc.write_input(atoms)                 # Write input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ground State Energy Test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo\n",
    "from ase.calculators.vasp import Vasp\n",
    "import shutil\n",
    "import os\n",
    "\n",
    "symbol_infos = [\n",
    "    SymbolInfo('Mo', 'bcc', 3.163),\n",
    "    SymbolInfo('Nb', 'bcc', 3.307),\n",
    "    SymbolInfo('Ta', 'bcc', 3.321),\n",
    "    SymbolInfo('V', 'bcc', 2.997),\n",
    "    SymbolInfo('W', 'bcc', 3.185)\n",
    "]\n",
    "\n",
    "for symbol_info in symbol_infos:\n",
    "    atoms = symbol_info.create_bulk_atoms()\n",
    "    dir=f'Your_Path/{symbol_info.formula}_coh'\n",
    "    if os.path.exists(dir):\n",
    "        shutil.rmtree(dir)\n",
    "    calc = Vasp(directory= dir,  \n",
    "            encut=500,                  # ENCUT\n",
    "            ediff=1e-6,                 # EDIFF\n",
    "            kspacing=0.15,              # KSPACING\n",
    "            sigma=0.1,                  # SIGMA\n",
    "            ismear=1,                   # ISMEAR\n",
    "            ibrion=2,                   # IBRION\n",
    "            isif=3,                     # ISIF\n",
    "            nsw=200,                    # NSW\n",
    "            gga='PE',                   # GGA\n",
    "            prec='Accurate',            # PREC\n",
    "            lasph=True,                 # LASPH\n",
    "            nelmin=4,                   # NELMIN\n",
    "            nelm=200,                   # NELM\n",
    "            kpar=4,                     # KPAR\n",
    "            lwave=False)                # LWAVE\n",
    "    calc.write_input(atoms)             # Write input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mono_Vacancy Calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo\n",
    "from wizard.calculator import MaterialCalculator\n",
    "from wizard.frames import MultiMol\n",
    "from wizard.io import read_xyz\n",
    "from calorine.calculators import CPUNEP\n",
    "\n",
    "symbol_infos = [\n",
    "    SymbolInfo('Mo', 'bcc', 3.163),\n",
    "    SymbolInfo('Nb', 'bcc', 3.308),\n",
    "    SymbolInfo('Ta', 'bcc', 3.321),\n",
    "    SymbolInfo('V', 'bcc', 2.997),\n",
    "    SymbolInfo('W', 'bcc', 3.185)\n",
    "]\n",
    "for symbol_info in symbol_infos:\n",
    "    atoms = symbol_info.create_bulk_atoms()\n",
    "    calc = CPUNEP('nep.txt')\n",
    "    material_calculator = MaterialCalculator(atoms, calc, symbol_info.formula, symbol_info.structure)\n",
    "    material_calculator.formation_energy_vacancy(supercell=(3,3,3))\n",
    "\n",
    "frames = read_xyz('MaterialProperties.xyz')\n",
    "for i in range(len(frames)):\n",
    "    atoms = frames[i]\n",
    "    dir=f'Your_Path/vacancies/{i}'\n",
    "    calc = Vasp(directory= dir,  \n",
    "            encut=500,                  # ENCUT\n",
    "            ediff=1e-6,                 # EDIFF\n",
    "            kspacing=0.15,              # KSPACING\n",
    "            sigma=0.1,                  # SIGMA\n",
    "            ismear=1,                   # ISMEAR\n",
    "            ibrion=2,                   # IBRION\n",
    "            isif=3,                     # ISIF\n",
    "            nsw=200,                    # NSW\n",
    "            gga='PE',                   # GGA\n",
    "            prec='Accurate',            # PREC\n",
    "            lasph=True,                 # LASPH\n",
    "            nelmin=4,                   # NELMIN\n",
    "            nelm=200,                   # NELM\n",
    "            kpar=4,                     # KPAR\n",
    "            lwave=False)                # LWAVE\n",
    "    atoms.calc = calc\n",
    "    calc.write_input(atoms)             # Write input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Self-Interstitial Calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.atoms import SymbolInfo\n",
    "from wizard.calculator import MaterialCalculator\n",
    "from wizard.frames import MultiMol\n",
    "from wizard.io import read_xyz\n",
    "from calorine.calculators import CPUNEP\n",
    "\n",
    "symbol_infos = [\n",
    "    SymbolInfo('Mo', 'bcc', 3.163),\n",
    "    SymbolInfo('Nb', 'bcc', 3.308),\n",
    "    SymbolInfo('Ta', 'bcc', 3.321),\n",
    "    SymbolInfo('V', 'bcc', 2.997),\n",
    "    SymbolInfo('W', 'bcc', 3.185)\n",
    "]\n",
    "sia_vectors = [(1,1,1),(1,0,0),(1,1,0)]\n",
    "for symbol_info in symbol_infos:\n",
    "    atoms = symbol_info.create_bulk_atoms()\n",
    "    calc = CPUNEP('nep.txt')\n",
    "    material_calculator = MaterialCalculator(atoms, calc, symbol_info.formula, symbol_info.structure)\n",
    "    for vector in sia_vectors:\n",
    "        material_calculator.formation_energy_sia(vector, supercell=(4,4,4))\n",
    "\n",
    "frames = read_xyz('MaterialProperties.xyz')\n",
    "for i in range(len(frames)):\n",
    "    atoms = frames[i]\n",
    "    dir=f'Your_Path/sias/{i}'\n",
    "    calc = Vasp(directory= dir,  \n",
    "            encut=500,                  # ENCUT\n",
    "            ediff=1e-6,                 # EDIFF\n",
    "            kspacing=0.15,              # KSPACING\n",
    "            sigma=0.1,                  # SIGMA\n",
    "            ismear=1,                   # ISMEAR\n",
    "            ibrion=2,                   # IBRION\n",
    "            isif=3,                     # ISIF\n",
    "            nsw=200,                    # NSW\n",
    "            gga='PE',                   # GGA\n",
    "            prec='Accurate',            # PREC\n",
    "            lasph=True,                 # LASPH\n",
    "            nelmin=4,                   # NELMIN\n",
    "            nelm=200,                   # NELM\n",
    "            kpar=4,                     # KPAR\n",
    "            lwave=False)                # LWAVE\n",
    "    atoms.calc = calc\n",
    "    calc.write_input(atoms)             # Write input files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading VASP Results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "# -----------------------------------------------------------------------------\n",
    "# The following code is sourced from\n",
    "# GPUMD:\n",
    "# https://github.com/brucefan1983/GPUMD\n",
    "#\n",
    "# Copyright (c) 2022 wang laosi, licensed under MIT License\n",
    "# -----------------------------------------------------------------------------\n",
    "\"\"\"\n",
    "\n",
    "import os, sys\n",
    "import numpy as np\n",
    "from ase.io import read, write\n",
    "from ase import Atoms, Atom\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def Convert_atoms(atom):\n",
    "    # 1 eV/Å^3 = 160.21766 GPa\n",
    "    xx,yy,zz,yz,xz,xy = -atom.calc.results['stress']*atom.get_volume() # *160.21766 \n",
    "    atom.info['virial'] = np.array([(xx, xy, xz), (xy, yy, yz), (xz, yz, zz)])\n",
    "    atom.calc.results['energy'] = atom.calc.results['free_energy']\n",
    "    del atom.calc.results['stress']\n",
    "    del atom.calc.results['free_energy']\n",
    "\n",
    "\n",
    "def find_vasprun(start_path='.'):\n",
    "    result = []\n",
    "    for root, dirs, files in os.walk(start_path):\n",
    "        if 'vasprun.xml' in files:\n",
    "            result.append(os.path.join(root, 'vasprun.xml'))\n",
    "    return result\n",
    "\n",
    "Path = 'dataset'\n",
    "file_list = find_vasprun(Path)\n",
    "\n",
    "cnum = 0     # total number of configuration\n",
    "atoms_list, err_list = [], []\n",
    "for dir_name in tqdm(file_list):\n",
    "    try:\n",
    "        atoms = read(dir_name.strip('\\n'), index=\":\")\n",
    "    except:\n",
    "        err_list.append(dir_name)\n",
    "        continue\n",
    "    for ai in range(len(atoms)):\n",
    "        Convert_atoms(atoms[ai])\n",
    "        atoms_list.append(atoms[ai])\n",
    "    cnum += len(atoms)\n",
    "\n",
    "write(Path + '/train.xyz', atoms_list, format='extxyz')\n",
    "print('The total number of configurations is: {} \\n'.format(cnum))\n",
    "\n",
    "if err_list:\n",
    "    print(\"The list of failed calculation files is as follows.\")\n",
    "    for err_dirname in err_list:\n",
    "        print(err_dirname)\n",
    "else:\n",
    "    print(\"All calculations are successful!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Merge Training Sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wizard.io import read_xyz\n",
    "from wizard.frames import MultiMol\n",
    "import glob\n",
    "\n",
    "files = glob.glob('dataset/*.xyz')\n",
    "train_set = []\n",
    "for file in files:\n",
    "    frames = read_xyz(file)\n",
    "    for atoms in frames:\n",
    "        atoms.info['config_type'] = file.split('/')[-1][:-4]\n",
    "    train_set.extend(frames)\n",
    "MultiMol(train_set).dump('dataset/train.xyz')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
